import os
import gluoncv
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
import mxnet as mx
from lib.nms.nms import gpu_nms_wrapper
from lib.config import get_default_config
import os,cv2,time
import numpy as np
from models.resnet_v1_101_fpn_dcn_rcnn import RFCN_Resnet
import mxnet.ndarray as nd
from lib.common import lsdir
import matplotlib.pyplot as plt
from lib.data.coco import COCODetection
from lib.transforms.bbox import bbox_image_pad_n
def main():
    device_id = 5
    nms_wrapper = gpu_nms_wrapper(thresh=0.3,device_id=device_id)
    net = RFCN_Resnet(pretrained="pretrained/fpn_dcn_coco-0000.params")
    net.collect_params().reset_ctx(ctx=mx.gpu(device_id))
    for img_name in lsdir(rootdir="../../dataset/coco/images/train2014",suffix=".jpg"):
        t0 = time.time()
        ori_img = cv2.imread(img_name)[:,:,:]
        fscale = 600.0 / ori_img.shape[1]
        img_resized = cv2.resize(ori_img, (0, 0), fx=fscale, fy=fscale)
        img_resized,_ = bbox_image_pad_n(n=32)(img_resized,None)
        print(img_resized.shape)
        img_float = img_resized.astype(np.float32)
        mean = np.array([103.06,115.90, 123.15])[np.newaxis,np.newaxis]
        img_float -= mean
        img_float = img_float[:,:,(2,1,0)]
        img_float = np.transpose(img_float,(2,0,1))
        data = mx.nd.array(img_float[np.newaxis],ctx = mx.gpu(device_id))

        im_info = nd.array([[data.shape[2],data.shape[3],3]],ctx=mx.gpu(device_id))
        rois, cls_score, bbox_pred = net(data,im_info)
        pred_bboxes,pred_scores,pred_labels = net.post_process(rois, bbox_pred, cls_score,nms_wrapper,data.shape,scale=fscale)
        print(time.time()-t0)
        gluoncv.utils.viz.plot_bbox(ori_img[:,:,::-1],bboxes = pred_bboxes,scores = pred_scores ,class_names=COCODetection.CLASSES,labels=pred_labels)
        plt.show()
if __name__ == '__main__':
    main()

